# ====================================
# Company: Zonge International, INC.
# Developer: Wanjie Feng
# Date: 10/18/2023
# Time: 9:54 AM
# Filename: ERT_Process.py
# IDE: PyCharm
# =====================================
import numpy as np
import matplotlib.pyplot as plt
import time
from matplotlib import cm
from scipy.interpolate import griddata
import matplotlib.colors as colors
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.patches import Rectangle
import io
import math

'''
    本程序模拟了手部UI读取数据、显示数据的过程。这里假设主机传到手部的文件格式为*.cvs格式，
    如后期数据格式有变，数据读取部分也要做相应改变。
    1. 获取硬件信息
    2. 项目参数
    3. 读取 *.csv 数据及部分系统状态信息
    4. 形成头文件
    5. 处理数据
    6. 显示数据
    7. 保存数据
    8. 重复3-7
    9. 采集结束
'''

# from Read_System import ReadSys
# from Read_Survey import ReadSurvey
# from Data_collection import ReadData
# from Compute_RhoIP import ComputeRhoIP
# from ERT_Plot import ERTPlot


# 排列图示
def ERTuiArray(ABMN, a1_dtype, a1_shape, Aspacing,type,eleNum):
    #print(ABMN)
    # print(type(ABMN))
    abmn = ABMN.encode('ISO-8859-1')
    # print(abmn)
    ABMN1 = np.fromstring(abmn, dtype=getattr(np, a1_dtype)).reshape(eval(a1_shape))
    print(ABMN1)
    ###3. 读取CVS数据及系统参数
    '''
        手部首先需要向主机发送数据采集的命令，规定采集的参数，然后主机才能返回数据。这里假设命令发送成果，数据成功上传。
        数据按每次叠加上传，例如我们设置命令为叠加5次，每次波形长度为8，那么手部会收到8个完整波形后再做相应的处理。
        这8个波形的数据在数据文件中是一个数据块，有相应的头文件与之对应。
    '''
    IPValue = {}
    # for AI in range(NS):
    #     print("------------------ai-------------------------")
    #     print(AI)
    print("------------------ai111111111111111111111111111111111-------------------------")
    print(print(ABMN1[0, :]))
    IPValue[0] = ComputeRhoIPArray(ABMN1[0, :], Aspacing,type)

    ###4. 画出图形
    return ERTuiArrayShow(IPValue, ABMN1, Aspacing,type,eleNum)
    # print("===================================开始进入ERT_PLOT============================================")


def ComputeRhoIPArray(ABMN, Aspacing,type):
    IPValues = {}

    if type == 0 or type ==1:
        TXA = (ABMN[0] - 1) * Aspacing  # Aspacing 点距 也是app 填写的
        TXB = (ABMN[1] - 1) * Aspacing
        RXM = (ABMN[2] - 1) * Aspacing
        RXN = (ABMN[3] - 1) * Aspacing
        IPValues['TXA'] = TXA
        IPValues['TXB'] = TXB
        IPValues['RXM'] = RXM
        IPValues['RXN'] = RXN
        IPValues['X'] = (RXM - TXB) / 2 + TXB  # 画图横坐标位置
        zz = (RXM - TXB) / 2 + Aspacing / 2  # 画图垂向坐标位置
        if zz<0:
            IPValues['Z'] = abs(zz)
        else:
            IPValues['Z'] = (RXM - TXB) / 2 + Aspacing / 2  # 画图垂向坐标位置
    else:
        TXA = (ABMN[0] - 1) * Aspacing  # Aspacing 点距 也是app 填写的
        TXB = (ABMN[3] - 1) * Aspacing
        RXM = (ABMN[1] - 1) * Aspacing
        RXN = (ABMN[2] - 1) * Aspacing
        IPValues['TXA'] = TXA
        IPValues['TXB'] = TXB
        IPValues['RXM'] = RXM
        IPValues['RXN'] = RXN
        IPValues['X'] = (RXM - RXN) / 2 + RXN  # 画图横坐标位置
        zz = (RXM - TXB) / 2 + Aspacing / 2  # 画图垂向坐标位置
        if zz<0:
            IPValues['Z'] = abs(zz)
        else:
            IPValues['Z'] = (RXM - TXB) / 2 + Aspacing / 2  # 画图垂向坐标位置


    print("IPValues['Z'] =")
    print(IPValues['Z'] )
    IPValues['n'] = ABMN[2] - ABMN[1]  # 隔离系数
    return IPValues


def ERTuiArrayShow(IPValues, ABMN, Aspacing,type,eleNum):
    '''
    本程序画出高密度电阻率法的等值线图和排列图
    :param IPValues:
    :return: NONE
'''
    ####这里画出两个图，一个是排列图，利用红色圆圈表示发射，用黑色圆表示接收
    # 最后一个IP的值
    # 获取相应的画图参数
    # Aspacing 点距

    NIP = len(IPValues)
    lastIP = IPValues[NIP - 1]
    print("------------- lastIP ------------" + str(lastIP))

    # TXA = lastIP['TXA']
    # TXB = lastIP['TXB']
    # RXM = lastIP['RXM']
    # RXN = lastIP['RXN']
    X = lastIP['X']
    Z = lastIP['Z']
    # rho = lastIP['Rho']
    # #####计算图的尺寸
    Min_X = np.min(ABMN[:, 0] - 1) * Aspacing
    # print(type)
    BE_MAX = np.max(ABMN[:,3])
    # print("BE_MAX == "+str(BE_MAX/2))
    #A_indx_sum 表示ABMN在温纳、斯伦贝谢、温纳-斯伦贝谢装置A=1时出现的次数作为深度的y轴的刻度
    A_indx_sum = np.sum(ABMN[:,0] == 1)+1
    #A_indx_sum 表示ABMN在温纳、斯伦贝谢、温纳-斯伦贝谢装置A=1时出现的次数作为深度的y轴的刻度
    A_indx_sum_ouji = np.sum(ABMN[:,3] == BE_MAX)+1
    print("A_indx_sum_ouji == "+str(A_indx_sum_ouji))
    if type == 0 or type == 1:
        Max_X = np.max(ABMN[:, 3])+2
        Min_Z = int(eleNum)
        Max_Z = int(eleNum)*Aspacing
    else:
        Max_X = np.max(ABMN[:, 3])+2
        Min_Z = int(eleNum)
        Max_Z = int(eleNum)*Aspacing
    # print("------------- Min_X ------------ " + str(Min_X))
    # print("------------- Max_X ------------ " + str(Max_X))
    # print("------------- Min_Z ------------ " + str(Min_Z))
    # print("------------- Max_Z ------------ " + str(Max_Z))
    # print("------------- Aspacing ------------ " + str(Aspacing))

    fig1, ax1 = plt.subplots(figsize=(15, 7.5))
    fig1.patch.set_alpha(0.)  # 或者其他值
    ax1.clear()  # Python清除上次画图的内容
    plt.xlim(0)  # 设置横向尺度
    if type == 0 or type == 1:#偶极-偶极，单极-偶极

        plt.ylim([0, int(A_indx_sum_ouji)])
    else:#温纳#斯伦贝谢#温纳-斯伦贝谢
        plt.ylim([0, int(A_indx_sum)])
    #plt.figure(figsize=(22, 10))
    # 设置纵向尺度
    #Radius = Aspacing / 20  # 圆表示画图点，圆的半径
    Radius = 0.1  # 圆表示画图点，圆的半径
    # 画出所有排列的地下记录点
    NS = np.size(ABMN, 0)

    for CI in range(NS):
        if type == 0 or type == 1:
            BE = int(ABMN[CI, 1])
            ME = int(ABMN[CI, 2])
            XE = int(ME - BE) / 2 + BE
            ZE = int(ME - BE)

            E_C = plt.Circle((XE, ZE), radius=Radius, fill=False, color='black', linewidth=1.0, \
                             clip_on=False)  # 画图点位置
        else:
            ME = ABMN[CI, 1]
            NE = ABMN[CI, 2]
            BE = ABMN[CI, 3]
            XE = (NE - ME) / 2 + ME
            ZE = (BE - NE)
            E_C = plt.Circle((XE, ZE), radius=Radius, fill=False, color='black', linewidth=1.0, \
                             clip_on=False)  # 画图点位置
        #print("E_C == "+str(E_C))
        # print("ABMN[0:0]： ",ABMN[0][0])
        plt.gca().add_artist(E_C)
    # 画出本次测量的地下记录的及发射接收位置
    TXA_C = plt.Circle((ABMN[0][0], 0), radius=Radius, fill=False, color='red', linewidth=2.0, \
                        clip_on=False)  # 发射A位置
    plt.gca().add_artist(TXA_C)
    TXB_C = plt.Circle((ABMN[0][1], 0), radius=Radius, fill=False, color='red', linewidth=2.0, \
                        clip_on=False)  # 发射B位置
    plt.gca().add_artist(TXB_C)
    RXM_C = plt.Circle((ABMN[0][2], 0), radius=Radius, fill=False, color='black', linewidth=2.0, \
                        clip_on=False)  # 接收M位置
    plt.gca().add_artist(RXM_C)
    RXN_C = plt.Circle((ABMN[0][3], 0), radius=Radius, fill=False, color='black', linewidth=2.0, \
                        clip_on=False)  # 接收N位置
    plt.gca().add_artist(RXN_C)
    if type == 0 or type == 1:
        BE = int(ABMN[0, 1])
        ME = int(ABMN[0, 2])
        XE = int(ME - BE) / 2 + BE
        ZE = int(ME - BE)
        # 电阻率画图点用实心圆表示
        Rho_C = plt.Circle((XE, ZE), radius=Radius, fill=True, color='blue', linewidth=2.0, \
                            clip_on=False)  # 拟断面图画图点
    else:
        ME = ABMN[0, 1]
        NE = ABMN[0, 2]
        BE = ABMN[0, 3]
        XE = (NE - ME) / 2 + ME
        ZE = (BE - NE)
        # 电阻率画图点用实心圆表示
        Rho_C = plt.Circle((XE, ZE), radius=Radius, fill=True, color='blue', linewidth=2.0, \
                           clip_on=False)  # 拟断面图画图点
    #print("Rho_C == "+str(Rho_C))
    plt.gca().add_artist(Rho_C)
    plt.gca().invert_yaxis()
    plt.gca().set_xlabel('Array Location(m)', size=14)
    plt.gca().set_ylabel('n-spacing', size=14)
    plt.gca().xaxis.set_label_position('top')
    plt.gca().xaxis.tick_top()
    plt.gca().set_aspect('equal')

    # abs_yticks_labs = [abs(num) for num in abs_yticks]

    if type == 0 or type == 1:
        abs_yticks=np.arange(0, A_indx_sum_ouji, 1)
        #print("abs_yticks = 0/1")
       # print(abs_yticks)
        plt.yticks(abs_yticks)
    else:
        abs_yticks=np.arange(0, A_indx_sum, 1)
       # print("abs_yticks = 3")
        #print(abs_yticks)
        plt.yticks(abs_yticks)
    plt.xticks(np.arange(Min_X, Max_X, 1))
    # 用隔离系数代替纵向坐标的标签
    labels = [item.get_text() for item in ax1.get_yticklabels()]
    print(labels)
    NLAB = len(labels)
    if type == 0 or type == 1:
        labs = np.arange(0, A_indx_sum_ouji, 1)
    else:
        labs = np.arange(0, A_indx_sum, 1)

    abs_labs = [abs(num) for num in labs]
   # print(abs_labs)
    ax1.set_yticklabels(abs_labs)
    # 隐藏图像的底部边界
    ax1.spines['bottom'].set_visible(False)
    # plt.title('高密度排列图示')
    ###这个部分本身是不需要的，但是在视电阻率图加了colorbar之后，两个图大小一样了，因此这个排列图也加了一个colorbar
    divider = make_axes_locatable(ax1)
    cax = divider.append_axes('right', size="2%", pad=0.1)
    colormap = cm.get_cmap('jet')
    sm = plt.cm.ScalarMappable(cmap=colormap)
    cbar = plt.colorbar(sm, cax=cax)
    cbar.remove()  # 不显示colorbar

    if NLAB == len(labs):
        ax1.set_yticklabels(labs)
    # 隐藏图像的底部边界
    ax1.spines['bottom'].set_visible(False)
    plt.title('高密度排列图示')
    ###这个部分本身是不需要的，但是在视电阻率图加了colorbar之后，两个图大小一样了，因此这个排列图也加了一个colorbar
    divider = make_axes_locatable(ax1)
    cax = divider.append_axes('right', size="2%", pad=0.1)
    colormap = cm.get_cmap('jet')
    sm = plt.cm.ScalarMappable(cmap=colormap)
    cbar = plt.colorbar(sm, cax=cax)
    cbar.remove()  # 不显示colorbar
    #####
    plt.draw()
    plt.pause(0.01)
    f = io.BytesIO()
    plt.savefig(f, format="png")
    print("___________________ERTPlot__________________22222")
    plt.close()
    return f.getvalue()


def ERTuiArrayABMN(num_electrodes0, Max_n_spacing0, Aspacing0, array_type0):
    # 开始采集之前，先计算排列，这里的参数不需要保存到数据文件里。
    # num_electrodes 电极总数（来自主机）
    # Max_n_spacing 最大隔离系数，若不设置，会根据电极总数计算，为 (num_electrodes - 3) （手部输入）
    # Aspacing 点距 （这个手部输入）
    # array_type 排列方式
    print("------------------ERTuiArray--------------------")
    num_electrodes = int(num_electrodes0)
    Max_n_spacing = int(Max_n_spacing0)
    Aspacing = Aspacing0
    array_type = array_type0
    print(num_electrodes)
    print(Max_n_spacing)
    print(Aspacing)
    print(array_type)
    # 最大隔离系数太大，系统自动调整
    if Max_n_spacing > num_electrodes - 3:
        print('最大隔离系数设置错误，隔离系数范围为' + str(1) + '到' + str(num_electrodes - 3))
        Max_n_spacing = num_electrodes - 3

    NS = np.sum(range(1, Max_n_spacing + 1)) + Max_n_spacing * (
            (num_electrodes - 3) - Max_n_spacing)  # 总排列数
    # print("-------------------------NS------------------------")
    # print(NS)
    #NS = num_electrodes*30
    F = 4
    ABMN = np.zeros((NS, F))  # ABMN电极排列
    ABMN1 = np.zeros((NS,F))                     #ABMN电极排列
    print(NS,F)
    A= None
    N= None
    M= None
    N = None
    c=0
    cc=0
    # 计算排列电极位置 （这里假设只有一个ADC接收电极，当有多个ADC时，此部分再修改)
    if array_type == 0:  # dipole-dipole
        for i in range(num_electrodes - 3):
            for j in range(i + 2, num_electrodes - 1):
                ABMN[c, 0] = i + 1  # A电极
                ABMN[c, 1] = i + 2  # B电极
                ABMN[c, 2] = j + 1  # M电极
                ABMN[c, 3] = j + 2  # N电极
                N = ABMN[c, 2] - ABMN[c, 1]
                if N > Max_n_spacing:
                    continue
                else:
                    c = c + 1
        print("dipole-dipole")
        print(ABMN)
    elif array_type == 1:  # pole-dipole
        for i in range(num_electrodes - 3):
            for j in range(i + 2, num_electrodes - 1):
                ABMN[c, 0] = 1  # A电极, 整个测量中不动
                ABMN[c, 1] = i + 2  # B电极
                ABMN[c, 2] = j + 1  # M电极
                ABMN[c, 3] = j + 2  # N电极
                N = ABMN[c, 2] - ABMN[c, 1]
                if N > Max_n_spacing:
                    continue
                else:
                    c = c + 1
    elif array_type == 2:#温纳 隔离系数 规律 1 2 3 1 2 3

        combinations = []
        for n in range(1, Max_n_spacing + 1):
            for a in range(1, num_electrodes - 3 * n + 1):
                m = a + n
                no = m + n
                b = no + n
                if b <= num_electrodes:
                    combinations.append((float(a), float(m), float(no), float(b)))
        ABMN = np.array(combinations)
        print("温纳")
        print(ABMN)
    elif array_type == 3:#斯伦贝谢
        n_min = 1
        combinations = []
        for n in range(n_min, Max_n_spacing + 1):
            # 计算AB和MN的半间距
            ab_half = n * Aspacing
            mn_half = Aspacing

            # 计算电极跨度
            span = int(ab_half / Aspacing + mn_half / Aspacing)  # 总跨度为AB/2 + MN/2对应的电极间隔数

            # 遍历所有可能的起始位置
            for start in range(1, num_electrodes - span + 1):
                A = start
                B = start + 2 * n + 1  # AB间距为2n*a，电极间隔数为2n
                M = start + n     # M位于A和B中间
                N = M + 1          # MN间距为2a，电极间隔数为2

                # 检查B是否超出电极总数
                if B > num_electrodes:
                    continue
                combinations.append([float(A), float(M), float(N), float(B)])
        ABMN = np.array(combinations)
        print("斯伦贝谢")
        print(ABMN)
    elif array_type == 4:#温纳-斯伦贝谢
        print("温纳-斯伦贝谢")
        combinations = []
        for n in range(1, Max_n_spacing + 1):
            # 计算当前n值下允许的最大起始电极位置
            max_start = num_electrodes - (2 * n + 1)
            if max_start < 1:
                continue  # 跳过无法生成任何组合的n值

            # 生成所有可能的起始位置组合
            for i in range(1, max_start + 1):
                A = i
                B = i + 2 * n + 1  # 供电电极B的位置
                M = A + n          # 测量电极M的位置
                N = M + 1          # 测量电极N的位置（与M相邻）
                combinations.append([float(A), float(M), float(N), float(B)])
        ABMN = np.array(combinations)

    dtype = str(ABMN.dtype)
    shape = str(ABMN.shape)
    print(dtype)
    print(shape)
    print(str(ABMN.tostring()))
    print(len(ABMN))
    print(ABMN)
    return str(len(ABMN)) + "," + str(dtype) + "," + str(shape) + "," + str(ABMN.tostring(),
                                                                            encoding="ISO-8859-1")
